import pandas as pd
import numpy as np
from statsmodels.tsa.stattools import adfuller, kpss
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
import ta  # 安装: pip install ta


class DataCleaner:
    def __init__(self, input_path):
        self.df = pd.read_csv(input_path, parse_dates=['trade_date'])
        self._validate_input()

    def _validate_input(self):
        """输入数据校验"""
        required = ['trade_date', 'price', 'buy_elg_vol']
        missing = set(required) - set(self.df.columns)
        if missing:
            raise KeyError(f"缺失关键字段: {missing}")

    def _handle_stationarity(self):
        """价格序列平稳化处理"""
        price_series = self.df['price'].dropna()

        # ADF检验
        adf_result = adfuller(price_series)
        # KPSS检验
        kpss_stat, kpss_pval = kpss(price_series)

        if adf_result > 0.05 or kpss_pval < 0.05:
            self.df['price_ret'] = self.df['price'].pct_change()
            self.df.drop(columns=['price'], inplace=True)

    def _build_features(self):
        """量化特征工程"""
        # 资金流特征
        self.df['net_flow'] = self.df['buy_elg_vol'] - self.df['sell_elg_vol']
        for lag in [1, 3, 5]:
            self.df[f'flow_lag_{lag}'] = self.df['net_flow'].shift(lag)

        # 技术指标
        self.df['rsi'] = ta.momentum.RSIIndicator(
            close=self.df['price_ret'].dropna()
        ).rsi()
        self.df['macd'] = ta.trend.MACD(
            close=self.df['price_ret'].dropna()
        ).macd_diff()

        # 指数相关特征
        self.df['index_ret'] = self.df['index_price'].pct_change()
        self.df['vol_ratio'] = self.df['index_vol'] / self.df['buy_elg_vol']

    def clean(self):
        """清洗主流程"""
        # 缺失值处理
        num_cols = ['price', 'buy_elg_vol', 'index_price']
        imputer = IterativeImputer(max_iter=15, random_state=42)
        self.df[num_cols] = imputer.fit_transform(self.df[num_cols])

        # 平稳性处理
        self._handle_stationarity()

        # 特征工程
        self._build_features()

        # 保存结果
        self.df.dropna().to_csv("cleaned_data.csv", index=False)


if __name__ == "__main__":
    cleaner = DataCleaner("raw_data.csv")
    cleaner.clean()
